package com.example.springbootdemo.text;

import java.util.HashMap;
import java.util.Map;
import java.util.Arrays;
import java.util.List;

public class TFIDFSimilarity {

    /**
     * 计算两个字符串的 TF-IDF 余弦相似度。
     * 
     * @param text1 第一个文本字符串
     * @param text2 第二个文本字符串
     * @return TF-IDF 余弦相似度
     */
    public static double computeTFIDFCosineSimilarity(String text1, String text2) {
        // 假设这两个文本为整个语料库
        List<String> documents = Arrays.asList(text1, text2);
        
        // 计算每个文档的词频
        List<Map<String, Integer>> termFrequencies = Arrays.asList(getTermFrequency(text1), getTermFrequency(text2));
        
        // 计算逆文档频率
        Map<String, Double> idf = computeIDF(documents, termFrequencies);
        
        // 计算 TF-IDF 向量
        Map<String, Double> tfidf1 = computeTFIDF(termFrequencies.get(0), idf);
        Map<String, Double> tfidf2 = computeTFIDF(termFrequencies.get(1), idf);
        
        // 计算余弦相似度
        return computeCosine(tfidf1, tfidf2);
    }
    
    /**
     * 计算文本的词频。
     */
    private static Map<String, Integer> getTermFrequency(String text) {
        String[] tokens = text.toLowerCase().replaceAll("[^a-z0-9\\s]", "").split("\\s+");
        Map<String, Integer> freqMap = new HashMap<>();
        for (String token : tokens) {
            if (!token.isEmpty()) {
                freqMap.put(token, freqMap.getOrDefault(token, 0) + 1);
            }
        }
        return freqMap;
    }
    
    /**
     * 计算逆文档频率（IDF）。
     */
    private static Map<String, Double> computeIDF(List<String> documents, List<Map<String, Integer>> termFrequencies) {
        Map<String, Double> idfMap = new HashMap<>();
        int totalDocuments = documents.size();
        
        // 统计每个词出现的文档数
        Map<String, Integer> docCount = new HashMap<>();
        for (Map<String, Integer> tf : termFrequencies) {
            for (String term : tf.keySet()) {
                docCount.put(term, docCount.getOrDefault(term, 0) + 1);
            }
        }
        
        // 计算 IDF
        for (String term : docCount.keySet()) {
            idfMap.put(term, Math.log((double) totalDocuments / (1 + docCount.get(term))));
        }
        
        return idfMap;
    }
    
    /**
     * 计算 TF-IDF 向量。
     */
    private static Map<String, Double> computeTFIDF(Map<String, Integer> termFrequency, Map<String, Double> idf) {
        Map<String, Double> tfidf = new HashMap<>();
        for (Map.Entry<String, Integer> entry : termFrequency.entrySet()) {
            String term = entry.getKey();
            int freq = entry.getValue();
            double tf = freq; // 简单的词频，未归一化
            double tf_idf = tf * idf.getOrDefault(term, 0.0);
            tfidf.put(term, tf_idf);
        }
        return tfidf;
    }
    
    /**
     * 计算两个 TF-IDF 向量的余弦相似度。
     */
    private static double computeCosine(Map<String, Double> vec1, Map<String, Double> vec2) {
        double dotProduct = 0.0;
        double normA = 0.0;
        double normB = 0.0;
        
        for (String term : vec1.keySet()) {
            double val1 = vec1.getOrDefault(term, 0.0);
            double val2 = vec2.getOrDefault(term, 0.0);
            dotProduct += val1 * val2;
            normA += Math.pow(val1, 2);
        }
        
        for (double val : vec2.values()) {
            normB += Math.pow(val, 2);
        }
        
        if (normA == 0.0 || normB == 0.0) {
            return 0.0;
        }
        
        return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
    }
    
    public static void main(String[] args) {
        String text1 = "我喜欢学习人工智能和机器学习。";
        String text2 = "机器学习和人工智能是我的兴趣所在。";
        
        double similarity = computeTFIDFCosineSimilarity(text1, text2);
        System.out.printf("TF-IDF 余弦相似度: %.4f\n", similarity);
    }
}
